import mnist
import experiment_runner
import numpy as np # 导入 numpy 用于生成范围
# ======== 日志重定向开始 ========

# --- 1. 加载 FashionMNIST 数据 ---
root_data,root_label,clients_data,clients_label,x_test,y_test = mnist.Load_FashionMNIST(q = 0.5, nclient = 100, root_dataset_size = 100, root_dataset_bias = 0.1)

# --- 2. 定义实验参数 ---
num_malicious_adaptive = 20
target_cosine_sim_values = np.arange(0.4, 0.6, 0.01)

# --- 3. 运行基准实验 (先运行这些) ---
#print("\n--- 运行基准: 无攻击 ---")
#server_no_attack = experiment_runner.run_no_attacks(root_data,root_label,clients_data,clients_label, 100, x_test, y_test)

#print("\n--- 运行基准: 仅服务器训练 ---")
#server_only_server = experiment_runner.run_only_server(root_data,root_label, x_test, y_test)

# --- 4. 运行自适应攻击训练循环 --

# 创建一个字典来存储每个 cos_sim 值训练得到的 server 对象
adaptive_servers = {}

for cos_sim in target_cosine_sim_values:
    # 使用 round 处理可能的浮点数精度问题，确保键和标签的一致性
    current_cos_sim = round(cos_sim, 3)
    print(f"\n--- 开始训练自适应攻击模型 (target_cosine_sim = {current_cos_sim:.2f}) ---")
    # 针对当前的 cos_sim 值运行完整的训练过程
    trained_server = experiment_runner.run_adaptive_attack(
        root_data, root_label,
        clients_data, clients_label,
        nclient=100,
        num_malicious=num_malicious_adaptive,
        target_cosine_sim=current_cos_sim, # 传递当前的余弦相似度值
        x_test=x_test, y_test=y_test
    )
    # 将训练完成的 server 对象存储到字典中
    adaptive_servers[current_cos_sim] = trained_server
    print(f"--- 完成训练 (cos_sim={current_cos_sim:.2f}) ---")

# --- 5. 最终评估阶段 ---
#server_only_server.evaluate(root_data, root_label, '[基准] 仅服务器 (训练集)')
#server_only_server.evaluate(x_test, y_test, '[基准] 仅服务器 (测试集)')
#server_no_attack.evaluate(x_test, y_test, '[基准] FLTrust 无攻击 (测试集)')

# 评估存储在字典中的自适应攻击模型
print("\n评估自适应攻击模型:")
# 按 cos_sim 值排序，以便按顺序输出结果
for cos_sim_key in sorted(adaptive_servers.keys()):
    server_to_evaluate = adaptive_servers[cos_sim_key] # 获取存储的 server 对象
    # 创建一个清晰的评估标签
    eval_label = f'[自适应攻击] FLTrust 测试集 ({num_malicious_adaptive}个恶意, cos={cos_sim_key:.2f})'
    # 调用 server 对象的 evaluate 方法进行评估
    server_to_evaluate.evaluate(x_test, y_test, eval_label)